import json
from PyQt6.QtWidgets import (QDialog, QVBoxLayout, QPushButton, QListWidget, 
                            QLineEdit, QHBoxLayout, QLabel, QMessageBox, QComboBox)
from PyQt6.QtCore import pyqtSignal, Qt
from styles import SETTINGS_BUTTON_STYLE, LABEL_STYLE, INPUT_STYLE, COMBOBOX_STYLE, DISABLED_INPUT_STYLE, BLUE_BUTTON_STYLE

class ModelManagerDialog(QDialog):
    model_updated = pyqtSignal()  # 用于通知主窗口模型列表已更新
    MAX_MODELS = 100  # 最大模型数量限制

    def __init__(self, openai_chat, parent=None):
        super().__init__(parent)
        self.openai_chat = openai_chat
        self.setWindowTitle("模型管理")
        self.setFixedSize(300, 400)
        self.temp_models = {}  # 临时存储模型配置
        self.changes_made = False  # 用于跟踪是否有未保存的更改
        self.init_ui()
        self.load_models()  # 从openai_chat加载当前配置到temp_models

    def init_ui(self):
        self.layout = QVBoxLayout(self)
        
        # 模型列表
        self.model_list = QListWidget()
        self.model_list.setSelectionMode(QListWidget.SelectionMode.ExtendedSelection)  # 支持多选
        self.model_list.setDragDropMode(QListWidget.DragDropMode.InternalMove)
        self.model_list.model().rowsMoved.connect(self.on_rows_moved)
        self.model_list.itemDoubleClicked.connect(self.select_model)
        self.model_list.setStyleSheet("""
            QListWidget {
                border: 1px solid #cccccc;
                border-radius: 4px;
                background-color: white;
            }
            QListWidget::item {
                padding: 5px;
                border-bottom: 1px solid #eeeeee;
            }
            QListWidget::item:selected {
                background-color: #007bff;
                color: white;
            }
            QListWidget::item:hover {
                background-color: #f5f5f5;
            }
            QListWidget::item:selected:hover {
                background-color: #0056b3;
                color: white;
            }
            QListWidget::item:selected:active {
                background-color: #004085;
                color: white;
            }
        """)
        self.layout.addWidget(self.model_list)
        
        # 添加标签显示模型数量
        self.count_label = QLabel()
        self.layout.addWidget(self.count_label)
        
        # 模型区域
        model_label = QLabel("模型名称:")
        model_label.setStyleSheet(LABEL_STYLE)
        self.layout.addWidget(model_label)
        
        self.model_input = QLineEdit()
        self.model_input.setPlaceholderText("输入模型名称")
        self.model_input.setStyleSheet(INPUT_STYLE + DISABLED_INPUT_STYLE)
        self.layout.addWidget(self.model_input)

        # API Key输入
        api_label = QLabel("API Key:")
        api_label.setStyleSheet(LABEL_STYLE)
        self.layout.addWidget(api_label)
        
        self.api_input = QLineEdit()
        self.api_input.setPlaceholderText("输入API Key")
        self.api_input.setStyleSheet(INPUT_STYLE)
        self.layout.addWidget(self.api_input)

        # Base URL输入
        base_url_label = QLabel("Base URL:")
        base_url_label.setStyleSheet(LABEL_STYLE)
        self.layout.addWidget(base_url_label)
        
        self.base_url_input = QLineEdit()
        self.base_url_input.setPlaceholderText("输入Base URL")
        self.base_url_input.setStyleSheet(INPUT_STYLE)
        self.layout.addWidget(self.base_url_input)

        # 按钮区域
        button_layout = QHBoxLayout()
        
        self.add_button = QPushButton("添加模型")
        self.add_button.setStyleSheet(SETTINGS_BUTTON_STYLE)
        self.add_button.clicked.connect(self.add_model)
        button_layout.addWidget(self.add_button)

        self.edit_save_button = QPushButton("编辑模型")
        self.edit_save_button.setStyleSheet(SETTINGS_BUTTON_STYLE)
        self.edit_save_button.clicked.connect(self.handle_edit_save)
        button_layout.addWidget(self.edit_save_button)

        # 修改为批量删除按钮
        self.delete_button = QPushButton("删除所选")
        self.delete_button.setStyleSheet(SETTINGS_BUTTON_STYLE)
        self.delete_button.clicked.connect(self.delete_selected_models)
        button_layout.addWidget(self.delete_button)

        self.save_button = QPushButton("保存更改")
        self.save_button.setStyleSheet(SETTINGS_BUTTON_STYLE)
        self.save_button.clicked.connect(self.save_changes)
        button_layout.addWidget(self.save_button)

        self.layout.addLayout(button_layout)

    def load_models(self):
        """加载模型列表"""
        self.model_list.clear()
        self.temp_models = self.openai_chat.get_all_models().copy()  # 复制当前模型配置到临时存储
        for model_name in self.temp_models:
            self.model_list.addItem(model_name)
        
        # 在加载完模型后更新数量显示
        self.update_count_label()

    def add_model(self):
        """添加新模型"""
        # 检查模型数量限制
        if len(self.temp_models) >= self.MAX_MODELS:
            QMessageBox.warning(self, "错误", f"模型数量已达到上限({self.MAX_MODELS})，请先删除一些模型！")
            return

        # 获取当前输入框的内容作为新模型
        model_name = self.model_input.text()
        api_key = self.api_input.text()
        base_url = self.base_url_input.text()

        # 验证输入
        if not all([model_name, api_key, base_url]):
            QMessageBox.warning(self, "错误", "请填写所有必要的信息！")
            return

        # 检查新名称是否已存在
        if model_name in self.temp_models:
            QMessageBox.warning(self, "错误", "模型名称已存在！")
            return

        # 直接添加为新模型
        self.temp_models[model_name] = {
            "api_key": api_key,
            "base_url": base_url,
            "model_name": model_name
        }
        self.model_list.addItem(model_name)
        self.changes_made = True
        
        # 清空输入并退出编辑状态
        self.clear_inputs()
        if hasattr(self, 'current_editing'):
            delattr(self, 'current_editing')
        
        # 保存更改
        self.save_changes()
        
        # 设置为当前模型并更新主界面
        self.openai_chat.models = self.temp_models.copy()
        self.openai_chat.set_current_model(model_name)  # 使用新方法
        self.model_updated.emit()
        self.accept()

    def handle_edit_save(self):
        """处理编辑/保存按钮点击"""
        current_item = self.model_list.currentItem()
        if not current_item:
            QMessageBox.warning(self, "错误", "请先选择一个模型！")
            return

        if self.model_input.text() == "":  # 输入框为空，说明是要编辑
            # 进入编辑状态
            model_name = current_item.text()
            model_data = self.temp_models.get(model_name)
            
            if model_data:
                self.model_input.setText(model_name)
                self.api_input.setText(model_data['api_key'])
                self.base_url_input.setText(model_data['base_url'])
                self.current_editing = model_name
                # 更改按钮文本和样式
                self.edit_save_button.setText("保存编辑")
                self.edit_save_button.setStyleSheet(BLUE_BUTTON_STYLE)
        else:  # 输入框有内容，说明是要保存
            new_model_name = self.model_input.text()
            new_api_key = self.api_input.text()
            new_base_url = self.base_url_input.text()

            # 验证输入
            if not all([new_model_name, new_api_key, new_base_url]):
                QMessageBox.warning(self, "错误", "请填写所有必要的信息！")
                return

            # 如果模型名称改变了，相当于添加新模型
            if new_model_name != self.current_editing:
                # 检查新名称是否已存在
                if new_model_name in self.temp_models:
                    QMessageBox.warning(self, "错误", "模型名称已存在！")
                    return
                # 删除原有模型
                del self.temp_models[self.current_editing]
                # 从列表中删除原有项
                for i in range(self.model_list.count()):
                    if self.model_list.item(i).text() == self.current_editing:
                        self.model_list.takeItem(i)
                        break
                # 添加新模型到列表
                self.model_list.addItem(new_model_name)
            
            # 更新或添加模型配置
            self.temp_models[new_model_name] = {
                "api_key": new_api_key,
                "base_url": new_base_url,
                "model_name": new_model_name
            }

            self.changes_made = True
            self.clear_inputs()
            delattr(self, 'current_editing')
            
            # 保存更改
            self.save_changes()
            
            # 设置为当前模型并更新主界面
            self.openai_chat.models = self.temp_models.copy()
            self.openai_chat.set_current_model(new_model_name)  # 使用新方法
            self.model_updated.emit()
            self.accept()

    def select_model(self, item):
        """双击选中模型并更新主界面"""
        model_name = item.text()
        # 设置为当前模型并更新主界面
        self.openai_chat.set_current_model(model_name)
        self.model_updated.emit()
        self.accept()  # 关闭模型管理窗口

    def delete_selected_models(self):
        """删除选中的模型"""
        selected_items = self.model_list.selectedItems()
        if not selected_items:
            QMessageBox.warning(self, "错误", "请先选择要删除的模型！")
            return

        reply = QMessageBox.question(
            self,
            "确认删除",
            f"确定要删除选中的 {len(selected_items)} 个模型吗？此操作不可恢复。",
            QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No,
            QMessageBox.StandardButton.No
        )

        if reply == QMessageBox.StandardButton.Yes:
            for item in selected_items:
                del self.temp_models[item.text()]
                self.model_list.takeItem(self.model_list.row(item))
            
            self.changes_made = True
            self.update_count_label()

    def clear_inputs(self):
        """清空输入框"""
        self.model_input.clear()
        self.api_input.clear()
        self.base_url_input.clear()
        # 恢复按钮文本和样式
        self.edit_save_button.setText("编辑模型")
        self.edit_save_button.setStyleSheet(SETTINGS_BUTTON_STYLE)

    def save_changes(self):
        """保存更改到配置文件并更openai_chat"""
        # 保存到文件
        with open('model_config.json', 'w', encoding='utf-8') as f:
            json.dump(self.temp_models, f, ensure_ascii=False, indent=4)
            
        # 更新openai_chat中的配置
        self.openai_chat.models = self.temp_models.copy()
        
        # 通知主界面更新
        self.model_updated.emit()
        
        self.changes_made = False
        QMessageBox.information(self, "保存成功", "模型配置已保存。")

    def closeEvent(self, event):
        """关闭窗口时检查是否有未保存的更改"""
        if self.changes_made:
            reply = QMessageBox.question(
                self, '未保存的更改',
                '您有未保存的更改，是否要保存？',
                QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No | QMessageBox.StandardButton.Cancel,
                QMessageBox.StandardButton.Cancel
            )

            if reply == QMessageBox.StandardButton.Yes:
                self.save_changes()
                event.accept()
            elif reply == QMessageBox.StandardButton.No:
                # 恢复原始数据
                self.temp_models = self.openai_chat.get_all_models().copy()
                self.model_updated.emit()
                event.accept()
            else:  # Cancel
                event.ignore()
        else:
            event.accept()

    def on_rows_moved(self, parent, start, end, destination, row):
        """处理拖拽完成事件"""
        self.changes_made = True
        
        # 更新 temp_models 中的顺序
        new_temp_models = {}
        for i in range(self.model_list.count()):
            model_name = self.model_list.item(i).text()
            new_temp_models[model_name] = self.temp_models[model_name]
        self.temp_models = new_temp_models
        
        # 临时更新主界面显示
        self.openai_chat.models = self.temp_models.copy()
        self.model_updated.emit()

    def update_count_label(self):
        """更新模型数量标签"""
        current_count = self.model_list.count()
        self.count_label.setText(f"当前模型数量: {current_count}/{self.MAX_MODELS}")
        self.count_label.setStyleSheet(
            "color: #ff0000;" if current_count >= self.MAX_MODELS else "color: #000000;"
        )

class ModelSelector(QComboBox):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setMinimumWidth(200)
        self.setFixedHeight(30)
        self.setStyleSheet(COMBOBOX_STYLE)

    def load_models(self, models, select_model=None):
        self.clear()  # 清除现有项目
        for model_name in models:
            self.addItem(model_name)
            
        if select_model and select_model in models:
            index = self.findText(select_model)
            if index >= 0:
                self.setCurrentIndex(index)

    def get_current_model(self):
        return self.currentText() 